{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4095e627-9bb2-44d7-82f1-58b27a1af1e0",
   "metadata": {},
   "source": [
    "# Using LLMs hosted on Vertex AI "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98cfbad3-ce56-4306-b996-7869fb9b007f",
   "metadata": {},
   "source": [
    "This guide teaches you how to use NeMo Guardrails with LLMs hosted on Vertex AI. It uses the [ABC Bot configuration](../../../../examples/bots/abc) and changes the model to `gemini-1.0-pro`.  \n",
    "\n",
    "This guide assumes you have configured and tested working with Vertex AI models. If not, refer to [this guide](../../advanced/vertexai-setup.md). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9cc0e5d657e75b33",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T14:52:37.023733Z",
     "start_time": "2024-03-15T14:52:36.842407Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Init: remove any existing configuration\n",
    "!rm -fr config\n",
    "\n",
    "# Get rid of the TOKENIZERS_PARALLELISM warning\n",
    "import warnings\n",
    "\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05fc110d-efb2-4e59-a962-6629c959f579",
   "metadata": {},
   "source": [
    "## Prerequisites\n",
    "\n",
    "You need to install the following Python libraries:"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "608db145d645cba",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "1. Install the `google-cloud-aiplatform` and `langchain-google-vertexai` packages:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fed8014-ecd8-4585-8781-63523e2cecf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install --quiet \"google-cloud-aiplatform>=1.38.0\" langchain-google-vertexai==0.1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36fbca4006c386d3",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "2. Set the `GOOGLE_APPLICATION_CREDENTIALS` environment variable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2b9d57c378a6fde1",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T14:52:39.121018Z",
     "start_time": "2024-03-15T14:52:39.004302Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "!export GOOGLE_APPLICATION_CREDENTIALS=$GOOGLE_APPLICATION_CREDENTIALS # Replace with your own key"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1322278e771b634",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "3. If you're running this inside a notebook, patch the AsyncIO loop."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "90b425e95950b75",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T14:52:39.126243Z",
     "start_time": "2024-03-15T14:52:39.121188Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dab99bd2-0568-49a4-85b9-2f8e2576c64b",
   "metadata": {},
   "source": [
    "## Configuration\n",
    "\n",
    "To get started, copy the ABC bot configuration into a subdirectory called `config`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "700c6d15-da11-4ec1-9146-6f76d1fd9215",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T14:52:39.253811Z",
     "start_time": "2024-03-15T14:52:39.126901Z"
    }
   },
   "outputs": [],
   "source": [
    "!cp -r ../../../../examples/bots/abc config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "71559518-5907-423b-a7c3-81daedd8a0ba",
   "metadata": {},
   "source": [
    "Update the `config/config.yml` file to use the `gemini-1.0-pro` model with the `vertexai` provider:\n",
    "\n",
    "```\n",
    "...\n",
    "\n",
    "models:\n",
    "  - type: main\n",
    "    engine: vertexai\n",
    "    model: gemini-1.0-pro\n",
    "\n",
    "...\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9c82b9b32f860286",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T14:52:39.259617Z",
     "start_time": "2024-03-15T14:52:39.254555Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Hide from documentation page.\n",
    "with open(\"config/config.yml\") as f:\n",
    "    content = f.read()\n",
    "\n",
    "content = content.replace(\n",
    "    \"\"\"\n",
    "  - type: main\n",
    "    engine: openai\n",
    "    model: gpt-3.5-turbo-instruct\"\"\",\n",
    "    \"\"\"\n",
    "  - type: main\n",
    "    engine: vertexai\n",
    "    model: gemini-1.0-pro\"\"\",\n",
    ")\n",
    "\n",
    "with open(\"config/config.yml\", \"w\") as f:\n",
    "    f.write(content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad931b8d621cfced",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Load the guardrails configuration:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba9b49e9-b550-4c1c-9fe2-7754d3358e43",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nemoguardrails import LLMRails, RailsConfig\n",
    "\n",
    "config = RailsConfig.from_path(\"./config\")\n",
    "rails = LLMRails(config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d986f0b2a43b1c9f",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "Test that it works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "2fc69196ab95b934",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T14:53:10.106244Z",
     "start_time": "2024-03-15T14:53:06.067506Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'role': 'assistant', 'content': \"I'm doing great! Thank you for asking. I'm here to help you with any questions you may have about the ABC Company.\"}\n"
     ]
    }
   ],
   "source": [
    "response = rails.generate(messages=[{\"role\": \"user\", \"content\": \"Hi! How are you?\"}])\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "979c8c65-d72e-4eac-b284-d26dc2609035",
   "metadata": {},
   "source": [
    "You can see that the bot responds correctly. To see in more detail what LLM calls have been made, you can use the `print_llm_calls_summary` method as follows: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "a3121315360899ce",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-15T14:53:13.141100Z",
     "start_time": "2024-03-15T14:53:13.132882Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Summary: 5 LLM call(s) took 3.99 seconds .\n",
      "\n",
      "1. Task `self_check_input` took 0.58 seconds .\n",
      "2. Task `generate_user_intent` took 1.19 seconds .\n",
      "3. Task `generate_next_steps` took 0.71 seconds .\n",
      "4. Task `generate_bot_message` took 0.88 seconds .\n",
      "5. Task `self_check_output` took 0.63 seconds .\n"
     ]
    }
   ],
   "source": [
    "info = rails.explain()\n",
    "info.print_llm_calls_summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc34d7aa3373b392",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Evaluation \n",
    "\n",
    "The `gemini-1.0-pro` and `text-bison` models have been evaluated for topical rails, and `gemini-1.0-pro` has also been evaluated as a self-checking model for hallucination and content moderation. Evaluation results can be found [here](../../../../docs/evaluation/README.md).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddc165e80bfdcd8f",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Conclusion\n",
    "\n",
    "In this guide, you learned how to connect a NeMo Guardrails configuration to a Vertex AI LLM model. This guide uses `gemini-1.0-pro`, however, you can connect any other model following the same steps. "
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
